#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Simple DES (S‑DES) — concise reference implementation (Python 3)
- 8-bit block, 10-bit key
- Encrypt/decrypt one byte or ASCII strings
- Optional brute-force key search
"""
from typing import List, Tuple

# --- Permutations & S-boxes (per assignment) ---
P10 = [3, 5, 2, 7, 4, 10, 1, 9, 8, 6]
P8  = [6, 3, 7, 4, 8, 5, 10, 9]
IP  = [2, 6, 3, 1, 4, 8, 5, 7]
IP_INV = [4, 1, 3, 5, 7, 2, 8, 6]
EP  = [4, 1, 2, 3, 2, 3, 4, 1]  # expansion/permutation
P4  = [2, 4, 3, 1]

# S-boxes: each is 4x4 of 2-bit strings (row by outer bits, col by inner bits)
S0 = [
    ["01","00","11","10"],
    ["11","10","01","00"],
    ["00","10","01","11"],
    ["11","01","00","10"],
]
S1 = [
    ["00","01","10","11"],
    ["10","11","01","00"],
    ["11","00","01","11"],
    ["10","01","00","11"],
]

# --- Bit helpers ---
def _permute(bits: str, table: List[int]) -> str:
    return ''.join(bits[i-1] for i in table)

def _ls(bits: str, n: int) -> str:
    n %= len(bits)
    return bits[n:] + bits[:n]

def _xor(a: str, b: str) -> str:
    return ''.join('1' if x!=y else '0' for x,y in zip(a,b))

def _sbox(inp4: str, box) -> str:
    r = int(inp4[0] + inp4[3], 2)  # row: outer bits
    c = int(inp4[1:3], 2)          # col: inner bits
    return box[r][c]

# --- Key schedule ---
def _subkeys(key10: int) -> Tuple[str,str]:
    k = f"{key10:010b}"
    k = _permute(k, P10)
    L, R = k[:5], k[5:]
    L, R = _ls(L, 1), _ls(R, 1)       # LS-1
    K1 = _permute(L+R, P8)
    L, R = _ls(L, 2), _ls(R, 2)       # (total 3)
    K2 = _permute(L+R, P8)
    return K1, K2

# --- Round function ---
def _fk(L: str, R: str, K: str) -> Tuple[str,str]:
    # E/P then XOR subkey
    t = _permute(R, EP)
    t = _xor(t, K)
    # S-boxes
    left, right = t[:4], t[4:]
    s = _sbox(left, S0) + _sbox(right, S1)
    s = _permute(s, P4)
    # XOR with left half
    return _xor(L, s), R

def encrypt_byte(plain: int, key10: int) -> int:
    """Encrypt an 8-bit integer with 10-bit key."""
    K1, K2 = _subkeys(key10)
    b = f"{plain & 0xFF:08b}"
    b = _permute(b, IP)
    L, R = b[:4], b[4:]
    L, R = _fk(L, R, K1)     # round 1
    L, R = R, L              # SWAP
    L, R = _fk(L, R, K2)     # round 2
    c = _permute(L+R, IP_INV)
    return int(c, 2)

def decrypt_byte(cipher: int, key10: int) -> int:
    """Decrypt an 8-bit integer with 10-bit key."""
    K1, K2 = _subkeys(key10)
    b = f"{cipher & 0xFF:08b}"
    b = _permute(b, IP)
    L, R = b[:4], b[4:]
    L, R = _fk(L, R, K2)     # reverse order
    L, R = R, L
    L, R = _fk(L, R, K1)
    p = _permute(L+R, IP_INV)
    return int(p, 2)

# --- String helpers (ASCII, 1 byte per char) ---
def encrypt_str(s: str, key10: int) -> bytes:
    return bytes(encrypt_byte(b, key10) for b in s.encode('latin1'))

def decrypt_bytes(b: bytes, key10: int) -> bytes:
    return bytes(decrypt_byte(x, key10) for x in b)

# --- Brute force (given one known P,C pair) ---
def brute_force_key(plain: int, cipher: int) -> List[int]:
    cand = []
    for k in range(1024):
        if encrypt_byte(plain, k) == cipher:
            cand.append(k)
    return cand

# --- Demo ---
if __name__ == "__main__":
    k = 0b1010000010   # 10-bit demo key (642)
    p = 0b11010111     # demo plaintext byte (215)
    c = encrypt_byte(p, k)
    d = decrypt_byte(c, k)
    print(f"key={k:010b} ({k})  plain={p:08b}  -> cipher={c:08b}")
    print(f"decrypt -> {d:08b}  (ok? {d==p})")
    # String demo
    msg = "Hello"
    enc = encrypt_str(msg, k)
    dec = decrypt_bytes(enc, k).decode('latin1', errors='ignore')
    print("enc bytes:", list(enc))
    print("dec text :", dec)
